#include "jit_compiler.h"
#include "../memory/memory_manager.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <time.h>
#include <sys/mman.h>
#include <unistd.h>
#include <errno.h>

// Platform-specific includes
#ifdef __x86_64__
#include <xmmintrin.h>
#endif

// Utility macros
#define HASH_SEED 0x12345678
#define FNV_OFFSET_BASIS 14695981039346656037ULL
#define FNV_PRIME 1099511628211ULL

// Internal helper functions
static uint64_t fnv1a_hash(const char* str);
static uint32_t next_power_of_2(uint32_t n);
static int make_memory_executable(void* ptr, size_t size);

// =============================================================================
// JIT Compiler Creation and Management
// =============================================================================

jit_compiler_t* jit_compiler_create(memory_manager_t* memory_mgr) {
    if (!memory_mgr) return NULL;
    
    jit_compiler_t* compiler = mm_alloc(memory_mgr, sizeof(jit_compiler_t), MEM_TYPE_QUERY_BUFFER);
    if (!compiler) return NULL;
    
    memset(compiler, 0, sizeof(jit_compiler_t));
    compiler->memory_mgr = memory_mgr;
    
    // Initialize expression cache
    compiler->cache_size = 1024;  // Initial cache size
    compiler->expression_cache = mm_alloc(memory_mgr, 
                                        compiler->cache_size * sizeof(jit_expression_t*),
                                        MEM_TYPE_QUERY_BUFFER);
    if (!compiler->expression_cache) {
        mm_free(memory_mgr, compiler);
        return NULL;
    }
    memset(compiler->expression_cache, 0, compiler->cache_size * sizeof(jit_expression_t*));
    
    // Allocate executable memory pool
    compiler->code_memory_size = JIT_CODE_CACHE_SIZE;
    compiler->code_memory = mmap(NULL, compiler->code_memory_size,
                               PROT_READ | PROT_WRITE | PROT_EXEC,
                               MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (compiler->code_memory == MAP_FAILED) {
        mm_free(memory_mgr, compiler->expression_cache);
        mm_free(memory_mgr, compiler);
        return NULL;
    }
    
    // Initialize synchronization primitives
    if (pthread_rwlock_init(&compiler->cache_lock, NULL) != 0 ||
        pthread_mutex_init(&compiler->code_mutex, NULL) != 0 ||
        pthread_mutex_init(&compiler->stats_mutex, NULL) != 0 ||
        pthread_mutex_init(&compiler->optimizer_mutex, NULL) != 0 ||
        pthread_cond_init(&compiler->optimizer_cond, NULL) != 0) {
        munmap(compiler->code_memory, compiler->code_memory_size);
        mm_free(memory_mgr, compiler->expression_cache);
        mm_free(memory_mgr, compiler);
        return NULL;
    }
    
    // Set default configuration
    compiler->optimization_level = 2;
    compiler->enable_profiling = true;
    compiler->enable_caching = true;
    compiler->compilation_threshold = JIT_COMPILATION_THRESHOLD;
    
    // Start background optimizer thread
    compiler->optimizer_running = true;
    if (pthread_create(&compiler->optimizer_thread, NULL, jit_optimizer_thread_func, compiler) != 0) {
        compiler->optimizer_running = false;
        // Continue without background optimizer
    }
    
    return compiler;
}

void jit_compiler_destroy(jit_compiler_t* compiler) {
    if (!compiler) return;
    
    // Stop background optimizer
    if (compiler->optimizer_running) {
        pthread_mutex_lock(&compiler->optimizer_mutex);
        compiler->optimizer_running = false;
        pthread_cond_signal(&compiler->optimizer_cond);
        pthread_mutex_unlock(&compiler->optimizer_mutex);
        pthread_join(compiler->optimizer_thread, NULL);
    }
    
    // Clean up expression cache
    pthread_rwlock_wrlock(&compiler->cache_lock);
    for (uint32_t i = 0; i < compiler->cache_size; i++) {
        jit_expression_t* expr = compiler->expression_cache[i];
        while (expr) {
            jit_expression_t* next = expr->next;
            jit_expression_destroy(expr);
            expr = next;
        }
    }
    pthread_rwlock_unlock(&compiler->cache_lock);
    
    // Free memory
    if (compiler->code_memory != MAP_FAILED) {
        munmap(compiler->code_memory, compiler->code_memory_size);
    }
    
    mm_free(compiler->memory_mgr, compiler->expression_cache);
    
    // Destroy synchronization primitives
    pthread_rwlock_destroy(&compiler->cache_lock);
    pthread_mutex_destroy(&compiler->code_mutex);
    pthread_mutex_destroy(&compiler->stats_mutex);
    pthread_mutex_destroy(&compiler->optimizer_mutex);
    pthread_cond_destroy(&compiler->optimizer_cond);
    
    mm_free(compiler->memory_mgr, compiler);
}

int jit_compiler_set_optimization_level(jit_compiler_t* compiler, uint32_t level) {
    if (!compiler || level >= JIT_OPTIMIZATION_LEVELS) return -1;
    
    compiler->optimization_level = level;
    return 0;
}

// =============================================================================
// Expression Compilation and Execution
// =============================================================================

jit_expression_t* jit_compile_expression(jit_compiler_t* compiler, const char* expression) {
    if (!compiler || !expression) return NULL;
    
    // Check cache first
    jit_expression_t* cached = jit_cache_lookup(compiler, expression);
    if (cached) {
        cached->execution_count++;
        return cached;
    }
    
    // Create new expression
    jit_expression_t* expr = mm_alloc(compiler->memory_mgr, sizeof(jit_expression_t), 
                                     MEM_TYPE_QUERY_BUFFER);
    if (!expr) return NULL;
    
    memset(expr, 0, sizeof(jit_expression_t));
    strncpy(expr->source, expression, sizeof(expr->source) - 1);
    expr->hash = jit_hash_expression(expression);
    expr->complexity_score = jit_estimate_expression_complexity(expression);
    
    // Determine if we should compile immediately
    expr->should_compile = (expr->complexity_score > 10) || 
                          (compiler->optimization_level >= 3);
    
    if (expr->should_compile) {
        // Create and compile function
        clock_t start_time = clock();
        
        // Create a simple function for the expression
        jit_type_t param_types[] = {JIT_TYPE_INT64, JIT_TYPE_DOUBLE}; // Example parameters
        expr->function = jit_function_create(compiler, "expr_func", JIT_TYPE_DOUBLE, 
                                           param_types, 2);
        
        if (expr->function) {
            // Build expression IR (simplified for demonstration)
            jit_context_t* ctx = jit_context_create(compiler);
            if (ctx) {
                // This is a placeholder - real implementation would parse expression
                // and generate appropriate IR instructions
                jit_value_t* param1 = jit_value_create_parameter(ctx, 0, JIT_TYPE_INT64);
                jit_value_t* param2 = jit_value_create_parameter(ctx, 1, JIT_TYPE_DOUBLE);
                jit_instruction_t* add_inst = jit_emit_binary_op(ctx, JIT_OP_ADD, param1, param2);
                if (add_inst) {
                    jit_emit_return(ctx, add_inst->dest);
                }
                
                // Compile the function
                if (jit_function_compile(expr->function) == 0) {
                    expr->is_hot = true;
                }
                
                jit_context_destroy(ctx);
            }
        }
        
        clock_t end_time = clock();
        expr->compilation_time = (uint64_t)((end_time - start_time) * 1000000 / CLOCKS_PER_SEC);
        
        pthread_mutex_lock(&compiler->stats_mutex);
        compiler->total_compilations++;
        compiler->total_compilation_time += expr->compilation_time;
        pthread_mutex_unlock(&compiler->stats_mutex);
    }
    
    // Add to cache
    jit_cache_insert(compiler, expr);
    
    return expr;
}

jit_value_t* jit_execute_expression(jit_expression_t* expr, jit_value_t** parameters) {
    if (!expr) return NULL;
    
    clock_t start_time = clock();
    jit_value_t* result = NULL;
    
    expr->execution_count++;
    
    // Check if we should compile now
    if (!expr->is_hot && expr->execution_count >= JIT_COMPILATION_THRESHOLD) {
        expr->should_compile = true;
        expr->is_hot = true;
        
        // Trigger compilation in background
        // (In a real implementation, this would queue compilation work)
    }
    
    if (expr->function && expr->function->is_compiled) {
        // Execute compiled code
        typedef double (*compiled_func_t)(int64_t, double);
        compiled_func_t func = (compiled_func_t)expr->function->compiled_function;
        
        if (func && parameters && parameters[0] && parameters[1]) {
            double compiled_result = func(parameters[0]->data.int_val, 
                                        parameters[1]->data.float_val);
            
            result = mm_alloc(expr->function->id ? 
                            ((jit_compiler_t*)expr->function)->memory_mgr : NULL, 
                            sizeof(jit_value_t), MEM_TYPE_TEMP_BUFFER);
            if (result) {
                result->type = JIT_TYPE_DOUBLE;
                result->data.float_val = compiled_result;
                result->is_constant = false;
                result->is_parameter = false;
            }
        }
    } else {
        // Interpret expression (fallback)
        result = mm_alloc(NULL, sizeof(jit_value_t), MEM_TYPE_TEMP_BUFFER);
        if (result) {
            result->type = JIT_TYPE_DOUBLE;
            result->data.float_val = 42.0; // Placeholder result
            result->is_constant = false;
            result->is_parameter = false;
        }
    }
    
    clock_t end_time = clock();
    uint64_t execution_time = (uint64_t)((end_time - start_time) * 1000000 / CLOCKS_PER_SEC);
    expr->total_time += execution_time;
    
    return result;
}

void jit_expression_destroy(jit_expression_t* expr) {
    if (!expr) return;
    
    if (expr->function) {
        jit_function_destroy(expr->function);
    }
    
    if (expr->result_cache.cached_result) {
        free(expr->result_cache.cached_result);
    }
    
    // Note: expr itself should be freed by the compiler when cleaning up cache
}

// =============================================================================
// Hash and Cache Management
// =============================================================================

uint64_t jit_hash_expression(const char* expression) {
    if (!expression) return 0;
    
    uint64_t hash = FNV_OFFSET_BASIS;
    const char* ptr = expression;
    
    while (*ptr) {
        hash ^= (uint64_t)*ptr;
        hash *= FNV_PRIME;
        ptr++;
    }
    
    return hash;
}

jit_expression_t* jit_cache_lookup(jit_compiler_t* compiler, const char* expression) {
    if (!compiler || !expression) return NULL;
    
    uint64_t hash = jit_hash_expression(expression);
    uint32_t index = hash % compiler->cache_size;
    
    pthread_rwlock_rdlock(&compiler->cache_lock);
    
    jit_expression_t* expr = compiler->expression_cache[index];
    while (expr) {
        if (expr->hash == hash && strcmp(expr->source, expression) == 0) {
            pthread_rwlock_unlock(&compiler->cache_lock);
            
            pthread_mutex_lock(&compiler->stats_mutex);
            compiler->cache_hits++;
            pthread_mutex_unlock(&compiler->stats_mutex);
            
            return expr;
        }
        expr = expr->next;
    }
    
    pthread_rwlock_unlock(&compiler->cache_lock);
    
    pthread_mutex_lock(&compiler->stats_mutex);
    compiler->cache_misses++;
    pthread_mutex_unlock(&compiler->stats_mutex);
    
    return NULL;
}

int jit_cache_insert(jit_compiler_t* compiler, jit_expression_t* expr) {
    if (!compiler || !expr) return -1;
    
    uint32_t index = expr->hash % compiler->cache_size;
    
    pthread_rwlock_wrlock(&compiler->cache_lock);
    
    // Insert at head of chain
    expr->next = compiler->expression_cache[index];
    if (compiler->expression_cache[index]) {
        compiler->expression_cache[index]->prev = expr;
    }
    compiler->expression_cache[index] = expr;
    compiler->cached_expressions++;
    
    pthread_rwlock_unlock(&compiler->cache_lock);
    
    return 0;
}

// =============================================================================
// Type System
// =============================================================================

jit_type_t jit_type_from_string(const char* type_name) {
    if (!type_name) return JIT_TYPE_UNKNOWN;
    
    if (strcmp(type_name, "int8") == 0) return JIT_TYPE_INT8;
    if (strcmp(type_name, "int16") == 0) return JIT_TYPE_INT16;
    if (strcmp(type_name, "int32") == 0) return JIT_TYPE_INT32;
    if (strcmp(type_name, "int64") == 0) return JIT_TYPE_INT64;
    if (strcmp(type_name, "uint8") == 0) return JIT_TYPE_UINT8;
    if (strcmp(type_name, "uint16") == 0) return JIT_TYPE_UINT16;
    if (strcmp(type_name, "uint32") == 0) return JIT_TYPE_UINT32;
    if (strcmp(type_name, "uint64") == 0) return JIT_TYPE_UINT64;
    if (strcmp(type_name, "float") == 0) return JIT_TYPE_FLOAT;
    if (strcmp(type_name, "double") == 0) return JIT_TYPE_DOUBLE;
    if (strcmp(type_name, "bool") == 0) return JIT_TYPE_BOOL;
    if (strcmp(type_name, "string") == 0) return JIT_TYPE_STRING;
    if (strcmp(type_name, "timestamp") == 0) return JIT_TYPE_TIMESTAMP;
    if (strcmp(type_name, "decimal") == 0) return JIT_TYPE_DECIMAL;
    if (strcmp(type_name, "null") == 0) return JIT_TYPE_NULL;
    
    return JIT_TYPE_UNKNOWN;
}

const char* jit_type_to_string(jit_type_t type) {
    switch (type) {
        case JIT_TYPE_INT8: return "int8";
        case JIT_TYPE_INT16: return "int16";
        case JIT_TYPE_INT32: return "int32";
        case JIT_TYPE_INT64: return "int64";
        case JIT_TYPE_UINT8: return "uint8";
        case JIT_TYPE_UINT16: return "uint16";
        case JIT_TYPE_UINT32: return "uint32";
        case JIT_TYPE_UINT64: return "uint64";
        case JIT_TYPE_FLOAT: return "float";
        case JIT_TYPE_DOUBLE: return "double";
        case JIT_TYPE_BOOL: return "bool";
        case JIT_TYPE_STRING: return "string";
        case JIT_TYPE_TIMESTAMP: return "timestamp";
        case JIT_TYPE_DECIMAL: return "decimal";
        case JIT_TYPE_NULL: return "null";
        case JIT_TYPE_UNKNOWN: return "unknown";
    }
    return "unknown";
}

size_t jit_type_size(jit_type_t type) {
    switch (type) {
        case JIT_TYPE_INT8:
        case JIT_TYPE_UINT8:
        case JIT_TYPE_BOOL:
            return 1;
        case JIT_TYPE_INT16:
        case JIT_TYPE_UINT16:
            return 2;
        case JIT_TYPE_INT32:
        case JIT_TYPE_UINT32:
        case JIT_TYPE_FLOAT:
            return 4;
        case JIT_TYPE_INT64:
        case JIT_TYPE_UINT64:
        case JIT_TYPE_DOUBLE:
        case JIT_TYPE_TIMESTAMP:
            return 8;
        case JIT_TYPE_STRING:
            return sizeof(char*);
        case JIT_TYPE_DECIMAL:
            return 16; // 128-bit decimal
        case JIT_TYPE_NULL:
            return 0;
        case JIT_TYPE_UNKNOWN:
            return 0;
    }
    return 0;
}

bool jit_type_is_numeric(jit_type_t type) {
    return type >= JIT_TYPE_INT8 && type <= JIT_TYPE_DECIMAL;
}

bool jit_type_is_integer(jit_type_t type) {
    return type >= JIT_TYPE_INT8 && type <= JIT_TYPE_UINT64;
}

jit_type_t jit_type_promote(jit_type_t type1, jit_type_t type2) {
    // Type promotion rules for arithmetic operations
    if (type1 == JIT_TYPE_DOUBLE || type2 == JIT_TYPE_DOUBLE) {
        return JIT_TYPE_DOUBLE;
    }
    if (type1 == JIT_TYPE_FLOAT || type2 == JIT_TYPE_FLOAT) {
        return JIT_TYPE_FLOAT;
    }
    if (type1 == JIT_TYPE_INT64 || type2 == JIT_TYPE_INT64) {
        return JIT_TYPE_INT64;
    }
    if (type1 == JIT_TYPE_UINT64 || type2 == JIT_TYPE_UINT64) {
        return JIT_TYPE_UINT64;
    }
    if (type1 == JIT_TYPE_INT32 || type2 == JIT_TYPE_INT32) {
        return JIT_TYPE_INT32;
    }
    if (type1 == JIT_TYPE_UINT32 || type2 == JIT_TYPE_UINT32) {
        return JIT_TYPE_UINT32;
    }
    
    return JIT_TYPE_INT32; // Default promotion
}

// =============================================================================
// Performance Monitoring and Statistics
// =============================================================================

void jit_compiler_print_stats(const jit_compiler_t* compiler) {
    if (!compiler) return;
    
    printf("JIT Compiler Statistics:\n");
    printf("========================\n");
    printf("Cached expressions: %u/%u\n", compiler->cached_expressions, compiler->cache_size);
    printf("Total compilations: %llu\n", (unsigned long long)compiler->total_compilations);
    printf("Total compilation time: %llu us\n", (unsigned long long)compiler->total_compilation_time);
    printf("Total executions: %llu\n", (unsigned long long)compiler->total_executions);
    printf("Cache hits: %llu\n", (unsigned long long)compiler->cache_hits);
    printf("Cache misses: %llu\n", (unsigned long long)compiler->cache_misses);
    printf("Cache hit ratio: %.2f%%\n", jit_compiler_get_cache_hit_ratio(compiler) * 100.0);
    printf("Average compilation time: %llu us\n", 
           (unsigned long long)jit_compiler_get_average_compilation_time(compiler));
    printf("Code memory used: %zu/%zu bytes (%.1f%%)\n", 
           compiler->code_memory_used, compiler->code_memory_size,
           (double)compiler->code_memory_used / compiler->code_memory_size * 100.0);
    printf("Optimization level: %u\n", compiler->optimization_level);
    printf("Profiling enabled: %s\n", compiler->enable_profiling ? "yes" : "no");
    printf("Caching enabled: %s\n", compiler->enable_caching ? "yes" : "no");
}

double jit_compiler_get_cache_hit_ratio(const jit_compiler_t* compiler) {
    if (!compiler) return 0.0;
    
    uint64_t total = compiler->cache_hits + compiler->cache_misses;
    if (total == 0) return 0.0;
    
    return (double)compiler->cache_hits / total;
}

uint64_t jit_compiler_get_average_compilation_time(const jit_compiler_t* compiler) {
    if (!compiler || compiler->total_compilations == 0) return 0;
    
    return compiler->total_compilation_time / compiler->total_compilations;
}

// =============================================================================
// Expression Complexity Estimation
// =============================================================================

uint32_t jit_estimate_expression_complexity(const char* expression) {
    if (!expression) return 0;
    
    uint32_t complexity = 0;
    const char* ptr = expression;
    
    while (*ptr) {
        switch (*ptr) {
            case '+':
            case '-':
                complexity += 1;
                break;
            case '*':
            case '/':
            case '%':
                complexity += 2;
                break;
            case '(':
            case ')':
                complexity += 1;
                break;
            case '=':
            case '<':
            case '>':
            case '!':
                complexity += 2;
                break;
            default:
                if (isalnum(*ptr)) {
                    complexity += 1;
                }
                break;
        }
        ptr++;
    }
    
    return complexity;
}

// =============================================================================
// Background Optimizer Thread
// =============================================================================

void* jit_optimizer_thread_func(void* arg) {
    jit_compiler_t* compiler = (jit_compiler_t*)arg;
    if (!compiler) return NULL;
    
    while (compiler->optimizer_running) {
        pthread_mutex_lock(&compiler->optimizer_mutex);
        
        // Wait for work or timeout
        struct timespec timeout;
        clock_gettime(CLOCK_REALTIME, &timeout);
        timeout.tv_sec += 5; // 5 second timeout
        
        int result = pthread_cond_timedwait(&compiler->optimizer_cond, 
                                          &compiler->optimizer_mutex, &timeout);
        
        if (!compiler->optimizer_running) {
            pthread_mutex_unlock(&compiler->optimizer_mutex);
            break;
        }
        
        pthread_mutex_unlock(&compiler->optimizer_mutex);
        
        // Perform adaptive optimization
        jit_adaptive_optimize(compiler);
        
        // Sleep to avoid busy waiting
        usleep(100000); // 100ms
    }
    
    return NULL;
}

int jit_adaptive_optimize(jit_compiler_t* compiler) {
    if (!compiler) return -1;
    
    // This is a placeholder for adaptive optimization logic
    // Real implementation would:
    // 1. Scan expression cache for hot expressions
    // 2. Recompile with higher optimization levels
    // 3. Perform profile-guided optimizations
    // 4. Adjust compilation thresholds based on performance
    
    return 0;
}

// =============================================================================
// Utility Functions
// =============================================================================

static uint64_t fnv1a_hash(const char* str) {
    uint64_t hash = FNV_OFFSET_BASIS;
    while (*str) {
        hash ^= (uint64_t)*str;
        hash *= FNV_PRIME;
        str++;
    }
    return hash;
}

static uint32_t next_power_of_2(uint32_t n) {
    if (n <= 1) return 1;
    n--;
    n |= n >> 1;
    n |= n >> 2;
    n |= n >> 4;
    n |= n >> 8;
    n |= n >> 16;
    return n + 1;
}

static int make_memory_executable(void* ptr, size_t size) {
    // This would use platform-specific APIs to make memory executable
    // On Linux/macOS, this might use mprotect()
    // On Windows, this might use VirtualProtect()
    
    #ifdef __unix__
    return mprotect(ptr, size, PROT_READ | PROT_EXEC);
    #else
    // Fallback - assume memory is already executable
    return 0;
    #endif
}